from PyQt6.QtWidgets import (QApplication, QMainWindow, QPushButton, 
                            QLabel, QWidget, QFileDialog, QLineEdit)
from PyQt6.QtGui import (QPixmap, QGuiApplication, QMouseEvent, QPaintEvent,
                        QPainter, QPen, QColor)
from PyQt6.QtCore import Qt, QRect


class ScreenshotTool:
    def __init__(self):
        """
        _summary_
        """
        from PyQt6.uic.load_ui import loadUi
        self.screenshot_pixmap: QPixmap = None  # 用于存储最终截图
        # 设置主窗口的 UI，包括“截图”按钮和“保存截图”按钮
        self.ui: QMainWindow = loadUi('screenshot.ui')
        self.screenshot_label: QLabel = self.ui.screenshot_label  # 显示截图的 QLabel
        self.capture_button: QPushButton = self.ui.capture_button  # “截图”按钮
        self.save_button: QPushButton = self.ui.save_button  # “保存截图”按钮
        self.checkBox = self.ui.checkBox  # 是否隐藏窗口的 QCheckBox
        self.lineEdit: QLineEdit = self.ui.lineEdit
        # 连接信号槽
        self.capture_button.clicked.connect(self.start_screenshot)
        self.save_button.clicked.connect(self.save_screenshot)

        self.screenshot_label.resizeEvent = self.resizeEvent
        self.selection_window = None

    def resizeEvent(self, event):
        """
        控件缩放事件。
        :param event: _description_
        :return :
        """
        if self.screenshot_pixmap is None:
            return  # 如果没有截图，直接返回
        label_size = self.screenshot_label.size()  # 获取label的当前尺寸
        pixmap_size = self.screenshot_pixmap.size()  # 获取原始截图的尺寸
        # 如果原始截图为空，直接返回
        if pixmap_size.isEmpty():
            return
        # 计算缩放比例
        width_ratio = (label_size.width()) / pixmap_size.width()
        height_ratio = (label_size.height()) / pixmap_size.height()
        scale_ratio = max(width_ratio, height_ratio)  # 选择较小的比例，以确保截图不会超过label的尺寸
        # 计算缩放后的尺寸
        scaled_width = int(pixmap_size.width() * scale_ratio)
        scaled_height = int(pixmap_size.height() * scale_ratio)
        # 缩放pixmap
        scaled_pixmap: QPixmap = self.screenshot_pixmap.scaled(
            scaled_width, scaled_height,
            Qt.AspectRatioMode.KeepAspectRatio,  # 保持比例
            Qt.TransformationMode.SmoothTransformation  # 平滑缩放
        )
        self.screenshot_label.setPixmap(scaled_pixmap)  # 设置到label
        print('label', label_size)
        print('pixmap', pixmap_size)
        print('scaled', scaled_width, scaled_height)

    def start_screenshot(self):
        """
        开始截图流程：隐藏主窗口，启动 SelectionWindow。
        :return :
        """
        if self.checkBox.isChecked():
            self.ui.hide()  # 隐藏主窗口
            import time
            time.sleep(0.2)  # 等待0.2s
        # 创建并显示 SelectionWindow
        self.screenshot_pixmap = None
        self.screenshot_label.clear()  # 清空截图内容
        self.screenshot_label.setText('截图在这里显示')
        self.selection_window = SelectionWindow(self)
        self.selection_window.showFullScreen()

    def set_screenshot(self, pixmap: QPixmap):
        """
        接收从 SelectionWindow 传递过来的截图，并显示在主窗口中。
        :params pixmap: 截图结果
        :return :
        """
        self.screenshot_pixmap = pixmap
        # 将截图缩放显示到 QLabel 中
        self.resizeEvent(None)
        self.save_button.setEnabled(True)  # 启用保存按钮
        self.ui.show()  # 显示主窗口
        self.selection_window = None

    def save_screenshot(self):
        """
        保存截图到本地文件。
        :return :
        """
        if self.screenshot_pixmap is not None:
            file_path, _ = QFileDialog.getSaveFileName(
                None,
                "保存截图",
                "",
                "PNG 图片 (*.png);;JPEG 图片 (*.jpg *.jpeg)"
            )
            if file_path:
                self.screenshot_pixmap.save(file_path)
                print(f"截图已保存到 {file_path}")
        else:
            print("没有可保存的截图")

class SelectionWindow(QWidget):

    def __init__(self, screenshotTool: ScreenshotTool):
        """
        _summary_
        :param screenshotTool: 截图工具窗口
        """
        super().__init__()
        self.screenshotTool = screenshotTool  # 父窗口（ScreenshotTool）
        self.start_point = None
        self.end_point = None
        self.is_selecting = False
        self.screenshot: QPixmap = None  # 屏幕截图
        # 设置窗口无边框、透明背景
        self.setWindowFlags(Qt.WindowType.FramelessWindowHint)
        self.setAttribute(Qt.WidgetAttribute.WA_TranslucentBackground)

        # 获取整个屏幕的截图
        screen = QGuiApplication.primaryScreen()
        dpiRatio: float = screen.devicePixelRatio()  # 变换到对应的物理坐标
        self.screenshot: QPixmap = screen.grabWindow(0)  # 抓取整个屏幕
        # 考虑DPI缩放
        self.screenshot = self.screenshot.scaled(
            int(self.screenshot.width() / dpiRatio),
            int(self.screenshot.height() / dpiRatio),
            Qt.AspectRatioMode.KeepAspectRatio,
            Qt.TransformationMode.SmoothTransformation
        )
        self.setGeometry(self.screenshot.rect())  # 设置窗口大小为屏幕大小

    def paintEvent(self, event: QPaintEvent):
        """
        绘制屏幕截图和选择区域。
        :param event: _description_
        :return :
        """
        painter = QPainter(self)
        # 1. 绘制整个屏幕截图作为背景
        # 绘制半透明黑色遮罩
        painter.setCompositionMode(QPainter.CompositionMode.CompositionMode_SourceOver)
        painter.fillRect(self.rect(), QColor(0, 0, 0, 128))

        if self.start_point and self.end_point:
            rect = QRect(self.start_point, self.end_point).normalized()
            # 清除选择区域的遮罩
            painter.setCompositionMode(QPainter.CompositionMode.CompositionMode_Clear)
            painter.fillRect(rect, Qt.GlobalColor.white)
            # 恢复默认绘制模式
            painter.setCompositionMode(QPainter.CompositionMode.CompositionMode_SourceOver)
            # 绘制选择框边框
            pen = QPen(Qt.GlobalColor.white, 2)
            painter.setPen(pen)
            painter.drawRect(rect)

    def mousePressEvent(self, event: QMouseEvent):
        """
        鼠标按下时记录起始点。
        :param event: _description_
        :return :
        """
        if event.button() == Qt.MouseButton.LeftButton:
            self.start_point = event.globalPosition().toPoint()
            self.end_point = event.globalPosition().toPoint()
            self.is_selecting = True
            self.update()

    def mouseMoveEvent(self, event: QMouseEvent):
        """
        鼠标移动时更新选择区域。
        :param event: _description_
        :return :
        """
        if self.is_selecting:
            self.end_point = event.globalPosition().toPoint()
            self.update()

    def mouseReleaseEvent(self, event: QMouseEvent):
        """
        鼠标松开时计算最终选择区域，裁剪截图并传递给父窗口。
        :param event: _description_
        :return :
        """
        if event.button() == Qt.MouseButton.LeftButton and self.is_selecting:
            self.is_selecting = False
            rect = QRect(self.start_point, self.end_point).normalized()
            # 从屏幕截图中裁剪出选择的区域
            x = rect.x()  # 起始点坐标 (QPoint)
            y = rect.y()
            w = rect.width()  # 区域大小 (QSize)
            h = rect.height()
            self.screenshotTool.lineEdit.setText(f'x:{x}    y:{y}    '
                                         f'w:{w}    h:{h}')
            cropped_pixmap = self.screenshot.copy(rect)
            # 将截图传递给父窗口
            self.screenshotTool.set_screenshot(cropped_pixmap)
            self.close()  # 关闭自己

if __name__ == "__main__":
    import sys
    app = QApplication(sys.argv)
    window = ScreenshotTool()
    window.ui.show()
    sys.exit(app.exec())
    